#coding:utf-8
import sys
import requests
from PyQt5.QtWidgets import (QApplication, QMainWindow, QVBoxLayout, QWidget,
                            QLabel, QLineEdit, QTextEdit, QPushButton,
                            QProgressBar, QHBoxLayout, QFileDialog, QSpinBox)
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from urllib.parse import urljoin, urlparse
from collections import OrderedDict, deque
import concurrent.futures

class PathProcessor:
    @staticmethod
    def extract_components(content):
        components = set()
        for line in content.split('\n'):
            line = line.strip()
            if not line:
                continue
                
            if line.startswith('http'):
                parsed = urlparse(line)
                path = parsed.path
            else:
                path = line
            
            parts = [p for p in path.split('/') if p and p != '.' and not p.startswith(':')]
            components.add(path.lstrip('/'))
            components.update(parts)
        
        return sorted(components, key=lambda x: x.count('/'))

class BruteWorker(QThread):
    update_progress = pyqtSignal(int, int)
    found_path = pyqtSignal(str, str)
    error_path = pyqtSignal(str, str)
    finished = pyqtSignal()
    status_update = pyqtSignal(str)
    paused = pyqtSignal(bool)

    def __init__(self, base_url, components, max_depth=3, max_threads=10):
        super().__init__()
        self.base_url = base_url.rstrip('/')
        self.components = components
        self.max_depth = max_depth
        self.max_threads = min(max_threads, 20)
        self.running = True
        self.pause_flag = False
        self.visited = set()
        self.valid_paths = set()
        self.queue = deque()
        self.current_depth = 1
        self.total_paths = 0
        self.completed_paths = 0

    def check_path(self, path):
        full_url = urljoin(self.base_url + '/', path)
        try:
            response = requests.head(full_url, timeout=5, allow_redirects=True)
            status = response.status_code
            
            if status in (200, 403):
                return (True, full_url, path)
            return (False, full_url, f"Status: {status}")
        except Exception as e:
            return (False, full_url, f"Error: {str(e)}")

    def pause(self):
        self.pause_flag = True
        self.paused.emit(True)

    def resume(self):
        self.pause_flag = False
        self.paused.emit(False)

    def run(self):
        try:
            # 第一层爆破
            self.status_update.emit(f"开始第{self.current_depth}层爆破...")
            for comp in sorted(self.components, key=lambda x: x.count('/')):
                if comp not in self.visited:
                    self.visited.add(comp)
                    self.queue.append((comp, 1))
            
            self.total_paths = len(self.queue)
            
            with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_threads) as executor:
                while self.queue and self.running and self.current_depth <= self.max_depth:
                    if self.pause_flag:
                        while self.pause_flag and self.running:
                            self.msleep(100)
                        if not self.running:
                            break
                    
                    current_path, depth = self.queue.popleft()
                    
                    future = executor.submit(self.check_path, current_path)
                    result = future.result()
                    
                    if result[0]:
                        self.found_path.emit(result[1], result[2])
                        self.valid_paths.add(current_path)
                        
                        # 如果找到有效路径且未达最大深度，准备下一层爆破
                        if depth < self.max_depth and '.' not in current_path.split('/')[-1]:
                            self.current_depth = depth + 1
                            self.status_update.emit(f"开始第{self.current_depth}层爆破...")
                            for comp in self.components:
                                new_path = f"{current_path}/{comp}"
                                if new_path not in self.visited:
                                    self.visited.add(new_path)
                                    self.queue.append((new_path, self.current_depth))
                                    self.total_paths += 1
                    else:
                        self.error_path.emit(result[1], result[2])
                    
                    self.completed_paths += 1
                    self.update_progress.emit(self.completed_paths, self.total_paths)
            
            self.status_update.emit("扫描完成")
        except Exception as e:
            self.status_update.emit(f"扫描出错: {str(e)}")
        finally:
            self.finished.emit()

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("信安西部-明镜高悬实验室 智能路由与API扫描系统 v1.0")
        self.setup_ui()
        self.components = []
        self.worker = None
        self.success_paths = []

    def setup_ui(self):
        central_widget = QWidget()
        layout = QVBoxLayout()

        url_layout = QHBoxLayout()
        url_layout.addWidget(QLabel("目标URL:"))
        self.url_input = QLineEdit()
        self.url_input.setPlaceholderText("http://example.com/")
        url_layout.addWidget(self.url_input)

        settings_layout = QHBoxLayout()
        
        depth_layout = QHBoxLayout()
        depth_layout.addWidget(QLabel("最大深度:"))
        self.depth_input = QSpinBox()
        self.depth_input.setRange(1, 10)
        self.depth_input.setValue(3)
        depth_layout.addWidget(self.depth_input)
        
        thread_layout = QHBoxLayout()
        thread_layout.addWidget(QLabel("线程数:"))
        self.thread_input = QSpinBox()
        self.thread_input.setRange(1, 20)
        self.thread_input.setValue(10)
        thread_layout.addWidget(self.thread_input)
        
        settings_layout.addLayout(depth_layout)
        settings_layout.addLayout(thread_layout)
        settings_layout.addStretch()

        self.components_display = QTextEdit()
        self.components_display.setReadOnly(True)

        btn_layout = QHBoxLayout()
        self.load_btn = QPushButton("加载TXT文件")
        self.load_btn.clicked.connect(self.load_txt_file)
        self.start_btn = QPushButton("开始扫描")
        self.start_btn.clicked.connect(self.start_scan)
        self.pause_btn = QPushButton("暂停")
        self.pause_btn.clicked.connect(self.toggle_pause)
        self.pause_btn.setEnabled(False)
        self.stop_btn = QPushButton("停止")
        self.stop_btn.clicked.connect(self.stop_scan)
        self.stop_btn.setEnabled(False)
        self.export_btn = QPushButton("导出结果")
        self.export_btn.clicked.connect(self.export_results)
        self.export_btn.setEnabled(False)
        btn_layout.addWidget(self.load_btn)
        btn_layout.addWidget(self.start_btn)
        btn_layout.addWidget(self.pause_btn)
        btn_layout.addWidget(self.stop_btn)
        btn_layout.addWidget(self.export_btn)

        self.progress = QProgressBar()
        self.status_label = QLabel("准备就绪")
        self.result_output = QTextEdit()
        self.result_output.setReadOnly(True)

        layout.addLayout(url_layout)
        layout.addLayout(settings_layout)
        layout.addWidget(QLabel("已加载路径组件:"))
        layout.addWidget(self.components_display)
        layout.addLayout(btn_layout)
        layout.addWidget(self.progress)
        layout.addWidget(self.status_label)
        layout.addWidget(QLabel("扫描结果:"))
        layout.addWidget(self.result_output)

        central_widget.setLayout(layout)
        self.setCentralWidget(central_widget)
        self.resize(1000, 700)

    def load_txt_file(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择TXT文件", "", "Text Files (*.txt)")
        if file_path:
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                    self.components = PathProcessor.extract_components(content)
                    self.components_display.setPlainText('\n'.join(self.components))
                    self.status_label.setText(f"已加载 {len(self.components)} 个路径组件")
            except Exception as e:
                self.status_label.setText(f"加载文件出错: {str(e)}")

    def start_scan(self):
        if not self.components:
            self.status_label.setText("请先加载有效的路径文件")
            return

        base_url = self.url_input.text().strip()
        if not base_url:
            self.status_label.setText("请输入目标URL")
            return

        max_depth = self.depth_input.value()
        max_threads = self.thread_input.value()

        self.result_output.clear()
        self.progress.setValue(0)
        self.start_btn.setEnabled(False)
        self.pause_btn.setEnabled(True)
        self.stop_btn.setEnabled(True)
        self.load_btn.setEnabled(False)
        self.export_btn.setEnabled(False)
        self.success_paths = []
        self.status_label.setText("扫描中...")

        self.worker = BruteWorker(base_url, self.components, max_depth, max_threads)
        self.worker.update_progress.connect(self.update_progress)
        self.worker.found_path.connect(self.add_result)
        self.worker.error_path.connect(self.add_error)
        self.worker.status_update.connect(self.status_label.setText)
        self.worker.finished.connect(self.scan_finished)
        self.worker.paused.connect(lambda x: self.pause_btn.setText("继续" if x else "暂停"))
        self.worker.start()

    def toggle_pause(self):
        if self.worker:
            if self.worker.pause_flag:
                self.worker.resume()
            else:
                self.worker.pause()

    def stop_scan(self):
        if self.worker:
            self.worker.running = False
            self.status_label.setText("正在停止...")

    def update_progress(self, current, total):
        self.progress.setMaximum(total)
        self.progress.setValue(current)

    def add_result(self, url, path):
        self.result_output.append(f"✅ 有效路径: {url} (路径: {path})")
        self.success_paths.append(f"{url} (路径: {path})")

    def add_error(self, url, msg):
        self.result_output.append(f"❌ {url} - {msg}")

    def scan_finished(self):
        self.start_btn.setEnabled(True)
        self.pause_btn.setEnabled(False)
        self.stop_btn.setEnabled(False)
        self.load_btn.setEnabled(True)
        self.export_btn.setEnabled(bool(self.success_paths))

    def export_results(self):
        if not self.success_paths:
            self.status_label.setText("没有可导出的结果")
            return

        file_path, _ = QFileDialog.getSaveFileName(self, "保存结果", "", "Text Files (*.txt)")
        if file_path:
            try:
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write("\n".join(self.success_paths))
                self.status_label.setText(f"结果已导出到 {file_path}")
            except Exception as e:
                self.status_label.setText(f"导出失败: {str(e)}")

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
